# -*- coding: utf-8 -*-
import sys
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QPushButton, QFileDialog, QVBoxLayout, QHBoxLayout, QMessageBox
from PyQt5.QtGui import QPixmap, QFont, QPalette, QBrush, QImage
from PyQt5.QtCore import Qt, QTimer
import torch
import torch.nn as nn
import torchvision.models as models
from PIL import Image
import torchvision.transforms as transforms
import os
import cv2
import numpy as np

# 定义资源路径函数
def resource_path(relative_path):
    """获取打包后的资源路径"""
    if hasattr(sys, '_MEIPASS'):
        return os.path.join(sys._MEIPASS, relative_path)
    return os.path.join(os.path.abspath("."), relative_path)

classes = ('粑粑柑', '白兰瓜', '白萝卜', '白心火龙果', '百香果', '菠萝', '菠萝莓', '菠萝蜜', '草莓', '车厘子', '番石榴-百', '番石榴-红', '佛手瓜', ' 甘蔗', '桂圆', '哈密瓜', '黑莓', '红苹果', '红心火龙果', '胡萝卜', '黄桃', '金桔', '橘子', '蓝莓', '梨', '李子', '荔枝', '莲雾', '榴莲', ' 芦柑', '芒果', '毛丹', '猕猴桃', '木瓜', '柠檬', '牛油果', '蟠桃', '枇杷', '葡萄-白', '葡萄-红', '脐橙', '青柠', '青苹果', '人参果', '桑葚', '沙果', '沙棘', '砂糖橘', '山楂', '山竹', '蛇皮果', '圣女果', '石榴', '柿子', '树莓', '水蜜桃', '酸角', '甜瓜-白', '甜瓜-金', '甜瓜-绿', '甜瓜-伊丽莎白', '沃柑', '无花果', '西瓜', '西红柿', '西梅', '西柚', '香蕉', '香橼', '杏', '血橙', '羊角蜜', '羊奶果', '杨梅', '杨桃', '腰 果', '椰子', '樱桃', '油桃', '柚子', '枣')

transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

DEVICE = torch.device("cpu")
# model = models.vgg16(weights=None)
# num_ftrs = model.classifier[6].in_features
# model.classifier[6] = nn.Linear(num_ftrs, len(classes))

# 改成外部模型路径
model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model.pth")
if not os.path.exists(model_path):
    raise FileNotFoundError(f"The model file {model_path} does not exist.")

model = models.vgg16(weights=None)
num_ftrs = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_ftrs, len(classes))
model.load_state_dict(torch.load(model_path, map_location=DEVICE))
model.to(DEVICE)
model.eval()

class FruitRecognitionApp(QWidget):
    def __init__(self):
        super().__init__()
        print("初始化函数开始")
        self.camera = cv2.VideoCapture(0)
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        self.is_camera_on = False
        self.initUI()
        print("初始化UI完成")

    def initUI(self):
        self.setWindowTitle('水果识别')
        self.setGeometry(100, 100, 1200, 800)
        palette = QPalette()

        # 使用 resource_path 加载背景图片
        # background_image_path = resource_path("Backdrop.jpg")
        # background_image = QPixmap(background_image_path)
        # if not background_image.isNull():
        #     palette.setBrush(QPalette.Background, QBrush(background_image))
        palette.setColor(QPalette.Background, Qt.white)  # 使用白色背景
        self.setPalette(palette)

        # 左侧显示原始图片或视频
        self.labelOriginal = QLabel(self)
        self.labelOriginal.setAlignment(Qt.AlignCenter)
        self.labelOriginal.setStyleSheet("border: 2px dashed black; padding: 10px; background-color: rgba(255, 255, 255, 150);")

        # 右侧显示结果图片
        self.labelResult = QLabel(self)
        self.labelResult.setAlignment(Qt.AlignCenter)
        self.labelResult.setStyleSheet("border: 2px dashed black; padding: 10px; background-color: rgba(255, 255, 255, 150);")

        # 显示预测结果的标签
        self.resultLabel = QLabel('预测类别: ')
        self.resultLabel.setFont(QFont('Arial', 18))

        # 按钮：上传图片、新图片上传、摄像头捕获
        self.uploadButton = QPushButton('上传图片')
        self.uploadButton.clicked.connect(self.uploadImage)
        self.newUploadButton = QPushButton('上传新图片')
        self.newUploadButton.clicked.connect(self.reset)
        self.cameraButton = QPushButton('打开摄像头')
        self.cameraButton.clicked.connect(self.toggle_camera)

        # 布局
        layout = QVBoxLayout()
        layout.addWidget(self.labelOriginal)
        layout.addWidget(self.labelResult)
        layout.addWidget(self.resultLabel)
        buttonLayout = QHBoxLayout()
        buttonLayout.addWidget(self.uploadButton)
        buttonLayout.addWidget(self.newUploadButton)
        buttonLayout.addWidget(self.cameraButton)
        layout.addLayout(buttonLayout)
        self.setLayout(layout)

        print("初始化界面")
        # 原有代码
        self.setWindowTitle('水果识别')
        # ...
        print("界面设置完毕")

    def uploadImage(self):
        options = QFileDialog.Options()
        fileName, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.jpg *.jpeg)", options=options)
        if fileName:
            self.displayImage(fileName)
            self.newUploadButton.setEnabled(True)
            # 直接调用识别逻辑
            self.predictImage(fileName)

    def reset(self):
        self.labelOriginal.clear()
        self.labelResult.clear()
        self.resultLabel.setText('预测类别: ')
        self.newUploadButton.setEnabled(False)

    def displayImage(self, imagePath):
        self.labelOriginal.setPixmap(QPixmap(imagePath).scaled(400, 400, Qt.KeepAspectRatio))

    def predictImage(self, image):
        try:
            print("开始预测")
            if isinstance(image, str):  # 如果是文件路径
                print(f"加载图片路径：{image}")
                image = Image.open(image)
            elif isinstance(image, QImage):  # 如果是 QImage 对象
                print("处理QImage对象")
                image = self.qimage_to_pil(image)
            else:
                print("未知图片类型")
            print(f"预处理前图像：{image}")
            # 图像预处理
            image_tensor = transform_test(image).unsqueeze(0).to(DEVICE)
            print(f"处理后tensor: {image_tensor.shape}")
            # 模型推理
            with torch.no_grad():
                outputs = model(image_tensor)
                print(f"模型输出: {outputs}")
                _, predicted = torch.max(outputs, 1)
                class_name = classes[predicted.item()]
                print(f"预测类别：{class_name}")
            self.resultLabel.setText(f'预测类别: {class_name}')
            # 显示图片逻辑
            # ...
        except Exception as e:
            print("预测出错:", e)
            import traceback; traceback.print_exc()

    def qimage_to_pil(self, qImg):
        """将 QImage 转换为 PIL.Image"""
        qImg = qImg.convertToFormat(QImage.Format_RGB888)
        width = qImg.width()
        height = qImg.height()
        ptr = qImg.bits()
        ptr.setsize(qImg.byteCount())
        arr = np.array(ptr).reshape(height, width, 3)  # 转换为 numpy 数组
        return Image.fromarray(arr)  # 转换为 PIL.Image

    def toggle_camera(self):
        """切换摄像头状态"""
        if self.is_camera_on:
            self.timer.stop()  # 停止定时器
            self.cameraButton.setText('打开摄像头')
            self.is_camera_on = False
        else:
            self.timer.start(30)  # 每30毫秒更新一帧
            self.cameraButton.setText('关闭摄像头')
            self.is_camera_on = True

    def update_frame(self):
        ret, frame = self.camera.read()
        if ret:
            try:
                print("捕获一帧")
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image_pil = Image.fromarray(frame_rgb)
                # 预测
                print("开始识别摄像头画面")
                with torch.no_grad():
                    image_tensor = transform_test(image_pil).unsqueeze(0).to(DEVICE)
                    outputs = model(image_tensor)
                    _, predicted = torch.max(outputs, 1)
                    class_name = classes[predicted.item()]
                    print(f"预测类别：{class_name}")
                # 绘制文字
                frame_with_text = cv2.putText(
                    frame_rgb,
                    f"Predicted: {class_name}",
                    (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1,
                    (0, 255, 0),
                    2
                )
                height, width, channel = frame_with_text.shape
                bytesPerLine = 3 * width
                qImg = QImage(frame_with_text.data, width, height, bytesPerLine, QImage.Format_RGB888)
                self.labelOriginal.setPixmap(QPixmap.fromImage(qImg).scaled(400, 400, Qt.KeepAspectRatio))
                # 更新预测文本
                self.resultLabel.setText(f'预测类别: {class_name}')
            except Exception as e:
                print("更新摄像头画面出错：", e)
                import traceback; traceback.print_exc()

if __name__ == '__main__':
    import traceback
    try:
        print("程序启动")
        app = QApplication(sys.argv)
        print("创建 QApplication 成功")
        ex = FruitRecognitionApp()
        print("创建窗口成功")
        ex.show()
        print("窗口显示成功")
        sys.exit(app.exec_())
    except Exception as e:
        print("程序运行时发生异常：", e)
        traceback.print_exc()